
from core.utils import agent_service
from core.graph import build_graph
from models.response import Response
from fastapi import Body
from typing import Dict
from loguru import logger

graph = build_graph()


async def data_fetch(
    sessionId: str = Body('', description='sessionId'),
    query: str = Body('', description='用户提出的问题')
) -> Dict:

    logger.info({
        "sessionId": sessionId,
        "question": query
    })

    try:
        # 创建sessionId， 如果没有的话。
        agent_service.call_agent(
            data={
                "agentId": '123456',
                "sessionId": sessionId,
                "question": query
            }
        )

        # 获取数据并整理。
        result = await graph.ainvoke(
            input={
                'sessionId': sessionId,
                'query': query,
                'userinfo': None,
                'schemas': [],
                'rewrites': [],
                'datasets': [],
                'filtered': []
            },
            stream_mode=["custom"])
        resp = Response(data=result[-1][-1])
    except Exception as e:
        resp = Response(msg='error', error=str(e))
    return resp.model_dump()
